from __future__ import with_statement

import sys, optparse, logging, csv

import common


logger = logging.getLogger("retraction_unfolding_filter")
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.INFO)


class RetractionPeak(object):

    def __init__(
        self, contour_length_m, start_force_n, rupture_force_n, extension_m):
        if rupture_force_n is not None:
            common.check_le(start_force_n, rupture_force_n)
        common.check_le(0, start_force_n)
        self.contour_length_m = contour_length_m
        self.start_force_n = start_force_n
        self.rupture_force_n = rupture_force_n
        common.check_ge(extension_m, 0)
        self.extension_m = extension_m

    def __repr__(self):
        return "RetractionPeak(%r, %r, %r, %r)" % (
            self.contour_length_m, self.start_force_n,
            self.rupture_force_n, self.extension_m)

    def __eq__(self, other):
        return type(other) is type(self) and self.__dict__ == other.__dict__

    def __ne__(self, other):
        return not self.__eq__(other)

    def with_force_factor(self, factor):
        return RetractionPeak(
            self.contour_length_m,
            self.start_force_n * factor, self.rupture_force_n * factor,
            self.extension_m)

    def with_rupture_force(self, rupture_force_n):
        return RetractionPeak(
            self.contour_length_m, self.start_force_n, rupture_force_n,
            self.extension_m)

    def with_force_multiplier(self, factor):
        return RetractionPeak(
            self.contour_length_m, self.start_force_n * factor,
            self.rupture_force_n * factor, self.extension_m)

class RetractionSegment(object):

    def __init__(
        self, retraction_velocity_m_per_s, peaks, noise_mean_n=0, noise_stddev_n=0):
        # Defaults for the old format.
        self.retraction_velocity_m_per_s = retraction_velocity_m_per_s
        common.check_increasing([peak.contour_length_m for peak in peaks])
        self.peaks = peaks
        self.noise_mean_n = noise_mean_n
        self.noise_stddev_n = noise_stddev_n

    def __len__(self):
        return len(self.peaks)

    def with_force_multiplier(self, factor):
        return RetractionSegment(
            self.retraction_velocity_m_per_s,
            [peak.with_force_multiplier(factor) for peak in self.peaks],
            self.noise_mean_n * factor, self.noise_stddev_n * factor)

    def detachment(self):
        return self.peaks[-1]

    def reversed_stretch_at_end(self, cl_range, detachment_strictly_largest):
        """
        Does NOT include the detachment peak in the returned stretch.
        """
        common.check_nondecreasing(cl_range)
        if len(self) < 2:
            return []
        a = []
        for i in range(len(self) - 2, -1, -1):
            cl_diff = self.peaks[i + 1].contour_length_m - self.peaks[i].contour_length_m
            common.check_gt(cl_diff, 0)
            if cl_range[0] <= cl_diff <= cl_range[1] and (
                self.peaks[i].rupture_force_n < self.peaks[-1].rupture_force_n or not detachment_strictly_largest):
                a.append(self.peaks[i])
            else:
                break
        return a

    def contour_length_rupture_force_extension_and_distance_of_peaks(
        self, spring_constant_n_per_m):
        """
        From contact point onwards.
        """
        return [(
                peak.contour_length_m, peak.rupture_force_n, peak.extension_m,
                max(0, peak.extension_m - peak.rupture_force_n / spring_constant_n_per_m))
                for peak in self.peaks]

    def __repr__(self):
        return "RetractionSegment(%r, %r, %r, %r)" % (
            self.retraction_velocity_m_per_s, self.peaks,
            self.noise_mean_n, self.noise_stddev_n)


class RetractionCurve(object):

    def __init__(self, path, spring_constant_n_per_m, a, b=None):
        self.path = path
        self.spring_constant_n_per_m = spring_constant_n_per_m
        if b is None:
            #common.check_ge(len(a), 1)
            self.segments = a
        else:
            # Support old style.
            self.segments = [RetractionSegment(a, b)]

    def __repr__(self):
        return "RetractionCurve(r'%s', %r, %r)" % (
            self.path, self.spring_constant_n_per_m, self.segments)

    def __len__(self):
        return len(self.segments)

    def with_force_multiplier(self, factor):
        return RetractionCurve(
            self.path, self.spring_constant_n_per_m,
            [segment.with_force_multiplier(factor) for segment in self.segments])

    def with_spring_constant(self, spring_constant_n_per_m):
        return RetractionCurve(self.path, spring_constant_n_per_m, self.segments)

    def last_detachment(self):
        return self.last_segment().detachment()

    def last_segment(self):
        return self.segments[-1]

    def contour_length_rupture_force_extension_and_distance_of_peaks(self):
        return self.last_segment().contour_length_rupture_force_extension_and_distance_of_peaks(
            self.spring_constant_n_per_m)

    def last_reversed_stretch_at_end(self, cl_range, detachment_strictly_largest):
        return self.last_segment().reversed_stretch_at_end(
            cl_range, detachment_strictly_largest)

    def last_retraction_velocity_m_per_s(self):
        return self.last_segment().retraction_velocity_m_per_s

    def last_retraction_noise_stddev_n(self):
        return self.last_segment().noise_stddev_n


class CurveStretch(object):
    
    def __init__(
        self, path, spring_constant_n_per_m, retraction_velocity_m_per_s,
        reversed_stretch_at_end, noise_stddev_n=0):
        self.path = path
        self.spring_constant_n_per_m = spring_constant_n_per_m
        self.retraction_velocity_m_per_s = retraction_velocity_m_per_s
        common.check_decreasing([
                peak.contour_length_m for peak in reversed_stretch_at_end
                if peak])
        self.reversed_stretch_at_end = reversed_stretch_at_end
        common.check_ge(noise_stddev_n, 0)
        self.noise_stddev_n = noise_stddev_n

    def __repr__(self):
        return "CurveStretch(r'%s', %r, %r, %r, %r)" % (
            self.path, self.spring_constant_n_per_m,
            self.retraction_velocity_m_per_s, self.reversed_stretch_at_end,
            self.noise_stddev_n)

    def __len__(self):
        return len(self.reversed_stretch_at_end)

    def reversed_forces_at_end(self):
        return [
            peak.rupture_force_n for peak in self.reversed_stretch_at_end if peak]

    def peak_for_peaks_ahead(self, peaks_ahead):
        return self.reversed_stretch_at_end[peaks_ahead]

    def unseen_peak_count(self):
        return sum([not peak for peak in self.reversed_stretch_at_end], 0)

    def with_unseen_peaks(self, n):
        return CurveStretch(
            self.path, self.spring_constant_n_per_m,
            self.retraction_velocity_m_per_s,
            [None] * n + self.reversed_stretch_at_end,
            self.noise_stddev_n)

    def with_spring_constant_n_over_m(self, k):
        return CurveStretch(
            self.path, k,
            self.retraction_velocity_m_per_s, [
                peak.with_force_factor(k / self.spring_constant_n_per_m)
                for peak in self.reversed_stretch_at_end], self.noise_stddev_n)

    def keep_only_last_peaks(self, n):
        return CurveStretch(
            self.path, self.spring_constant_n_per_m,
            self.retraction_velocity_m_per_s, self.reversed_stretch_at_end[:n],
            self.noise_stddev_n)

    @staticmethod
    def read_curve_stretches(curve_stretch_playlist):
        curve_stretches = []
        with open(curve_stretch_playlist, "rt") as f:
            for line in f:
                line = line.strip()
                if line:
                    curve_stretch = eval(line)
                    curve_stretches.append(curve_stretch)
        return curve_stretches

    @staticmethod
    def v2indices(curve_stretches):
        d = {}
        for i, curve_stretch in enumerate(curve_stretches):
            a = d.setdefault(
                curve_stretch.retraction_velocity_m_per_s, [])
            a.append(i)
        return d

    @staticmethod
    def kv2indices(curve_stretches):
        d = {}
        for i, curve_stretch in enumerate(curve_stretches):
            kv = curve_stretch.spring_constant_n_per_m * curve_stretch.retraction_velocity_m_per_s
            a = d.setdefault(kv, [])
            a.append(i)
        return d

    @staticmethod
    def forces_list(curve_stretches_list):
        return list(
            sum([
                    curve_stretch.reversed_forces_at_end()
                    for curve_stretch in curve_stretches], [])
            for curve_stretches in curve_stretches_list)

def add_suffix(
    curve_stretch_playlist_base, detachment_strictly_largest,
    minimum_detachment_force, min_interpeaks, cl_range):
    return  "%s_%s_%s_%d_%s_%s" % (
        curve_stretch_playlist_base,
        detachment_strictly_largest, minimum_detachment_force * 1e12,
        min_interpeaks, cl_range[0] * 1e9, cl_range[1] * 1e9)

def main(params):
    parser = optparse.OptionParser()
    parser.add_option("--retraction_playlist", help=".retractionPlaylist file") 
    parser.add_option("--cl_start_nm", type="float", help="lower bound on delta contour length")
    parser.add_option("--cl_end_nm", type="float", help="upper bound on delta contour length")
    parser.add_option(
        "--min_interpeaks", type="int",
        help="minimum number of peaks that, after unfolding, reveal the correct delta contour length")
    parser.add_option(
        "--max_interpeaks", type="int",
        help="")
    parser.add_option(
        "--detachment_strictly_largest", choices=("false", "False", "true", "True"),
        help="upper bound on delta contour length", default="True")
    parser.add_option("--minimum_detachment_force_pn")
    parser.add_option("--curve_stretch_playlist_base")
    parser.add_option("--peak_csv_base")
    
    options, args = parser.parse_args(params)
    if args:
        parser.print_help()
        return -1
    print(options)

    cl_range = (float(options.cl_start_nm) / 1e9, float(options.cl_end_nm) / 1e9)
    retraction_playlist = options.retraction_playlist
    curve_stretch_playlist_base = options.curve_stretch_playlist_base
    min_interpeaks = int(options.min_interpeaks)
    max_interpeaks = int(options.max_interpeaks)
    detachment_strictly_largest = {"false" : False, "true" : True}.get(options.detachment_strictly_largest.lower())
    minimum_detachment_force = float(options.minimum_detachment_force_pn) / 1e12

    curve_stretches = []
    with open(retraction_playlist, "rt") as f:
        for line in f:
            line = line.strip()
            if line:
                curve = eval(line)
                rupture_force_n = curve.last_detachment().rupture_force_n if len(curve) and len(curve.last_segment()) else 0
                # Allow None detachment force.
                if len(curve) and len(curve.last_segment()) and rupture_force_n >= minimum_detachment_force:
                    a = curve.last_reversed_stretch_at_end(
                        cl_range, detachment_strictly_largest)
                    if len(a) >= min_interpeaks:
                        curve_stretch = CurveStretch(
                            curve.path, curve.spring_constant_n_per_m,
                            curve.last_retraction_velocity_m_per_s(), a,
                            curve.last_retraction_noise_stddev_n())
                        curve_stretches.append(curve_stretch.keep_only_last_peaks(
                                max_interpeaks))
    # ${detachment_strictly_largest}_${minimum_detachment_force_pn}_${min_interpeaks}_${cl_start}_${cl_end}
    
    if curve_stretch_playlist_base:
        curve_stretch_playlist = add_suffix(
            curve_stretch_playlist_base, detachment_strictly_largest,
            minimum_detachment_force, min_interpeaks, cl_range
            ) + ".curveStretchPlaylist"
        with open(curve_stretch_playlist, "wt") as f:
            for curve_stretch in curve_stretches:
                f.write("%r\n" % curve_stretch)
        logger.info(
            "Wrote %d curve stretches totalling %d peaks in: %s",
            len(curve_stretches), sum(map(len, curve_stretches)),
            curve_stretch_playlist)

    peak_csv_base = options.peak_csv_base
    if peak_csv_base:
        peak_csv_path = add_suffix(
            peak_csv_base, detachment_strictly_largest,
            minimum_detachment_force, min_interpeaks, cl_range) + ".csv"
        logger.info("peak_csv_path: %s", peak_csv_path)
        with open(peak_csv_path, "wb") as out_file:
            writer = csv.writer(out_file)
            for i, curve_stretch in enumerate(curve_stretches):
                curve_stretch_row = [
                    i, curve_stretch.spring_constant_n_per_m,
                    curve_stretch.retraction_velocity_m_per_s]
                for peaks_ahead in range(len(curve_stretch)):
                    peak = curve_stretch.peak_for_peaks_ahead(peaks_ahead)
                    peak_row = curve_stretch_row + [
                        peaks_ahead, peak.contour_length_m, peak.rupture_force_n]
                    writer.writerow(peak_row)

if __name__ == "__main__":
    main(sys.argv[1:])
